import tqsdk.ta as ta
import numpy as np
import pandas as pd
import tqsdk.backtest


def count_feature(X: pd.DataFrame, target=None) -> pd.DataFrame:
    """
    利用天勤SDK计算技术因子,默认包括
    BBI 3 6 12 24, BIAS 12, CCI 14, DMA 10 50 10, EXPMA 12 28, KDJ 9 3 3, MACD DIFF, MACD MACD,
     MFI, MI, MTM, PRICEOSC, PSY, ROC
    :param target: 要计算的技术因子,dict类型 {"因子名" : [参数列表] }
    :param X: DataFrame对象,必须包含以下列
            open:  (K线起始时刻的最新价) 其实可有可无
            high:  (K线时间范围内的最高价)
            low:  (K线时间范围内的最低价)
            close:  (K线结束时刻的最新价)
            volume:  (K线时间范围内的成交量)
    :return: 添加了因子后的X
    """
    if target is None:
        target = {
            "BBI": [3, 6, 12, 24],  # 单位为价格的指标，与收盘价进行比较
            "BIAS": [12],  # 百分比值，表示超买或超卖情况
            "CCI": [14],  # 顺势指标，超买超卖情况，仅在+100与-100以外有指导意义
            "DMA": [10, 50, 10],  # ddd线向或上向下穿过ama线是交易信号，单位和价格相同
            # "EXPMA": [12, 20],
            "KDJ": [9, 3, 3],  # k,d 线为0~100 j为3k-2d 相互穿过产生交易信号
            "MACD": [12, 26, 9],  # 长短线差与其移动平均，中线买卖的指标参考 单位和价格相同
            "MFI": [14],  # 资金流量指数 0-100
            "MTM": [6, 10],
            # 当日收盘价 - N日前收盘价，单位与价格相同 主要观察MTM线相关与0的变化，同时与MTMMA相比可以看出超卖或超买趋势 https://zhuanlan.zhihu.com/p/503599662
            "PRICEOSC": [26, 12],  # -100~100 短均线减去以长均线再除以短均线得出的比例，表示偏离率，自上而下穿过0线和自下而上穿过零线产生交易信号

            "PSY": [12, 6],  # 上涨周期数占总周期数的百分比取值范围0~100
            "ROC": [12, 6],  # 当日减n日前再除以n日前，由持续运行在0轴下方或上方突破到另一方产生交易信号
            "WR": [4]  # 0~100 其中0~20是超买20-80为横盘80~100为超卖
        }

    feature_list = [getattr(ta, func)(X, *args) for func, args in target.items()]
    return pd.concat([X] + feature_list, axis=1).dropna()


def count_feature4numpy(batch_X: np.ndarray) -> np.ndarray:
    """
    为numpy类型的序列X计算技术因子
    :param batch_X: 前五列的数组分别为 "open", "high", 'low', "close", "volume"
    :return: 加上技术因子的array
    """
    res = []
    for i in range(batch_X.shape[0]):
        X = batch_X[i]
        X = pd.DataFrame(X, columns=["open",
                                     "high",
                                     'low',
                                     "close",
                                     "volume"] +
                                    [f"feature{i}" for i in range(X.shape[1] - 5)])

        res.append(np.array(count_feature(X)))

    return np.array(res)


def open0(df: pd.DataFrame, ):
    """
    计算开盘相对收盘价的占比 返回new_df[open0] = (df['open']/df['close'] - 1) *20 - 0.5
    :param df:输入数据
    :return:答案数据
    """
    df['open0'] = (df['open'] / df['close'] - 1) * 20 + 0.5
    return df


def high0(df: pd.DataFrame):
    """
    计算开盘相对收盘价的占比 返回new_df[clos0] = (df['high'] / df['close'] - 1) * 30
    :param df:输入数据
    :return:答案数据
    """
    df['high0'] = (df['high'] / df['close'] - 1) * 30
    return df


def low0(df: pd.DataFrame):
    """
    计算开盘相对收盘价的占比 返回new_df[clos0] = (1 - df['low'] / df['close']) * 30
    :param df:输入数据
    :return:答案数据
    """
    df['low0'] = (1 - df['low'] / df['close']) * 30
    return df


def standardization_cci(X: pd.DataFrame, k: float = 0.06) -> pd.DataFrame:
    """
    标准化 CCI 分别做中间值在100和中间值在-100的区间做sigmoid运算然后加和除以2映射到0-1区间
    standardization = 1 / (1 + np.exp(-k * (cci_series - 100))) + 1 / (1 + np.exp(-k * (cci_series + 100)))
    standardization = standardization / 2
    :param X:
    :param k:超参数k
    :return: 新增‘s_cci'列的X
    """
    assert 'cci' in X.columns.values
    cci = X['cci'].values
    ans = 1 / (1 + np.exp(-k * (cci - 100))) + 1 / (1 + np.exp(-k * (cci + 100)))
    X['s_cci'] = ans / 2
    return X


def standardization_macd(X: pd.DataFrame, k=10) -> pd.DataFrame:
    """
    在 https://arxiv.org/abs/1911.10107
    Deep Reinforcement Learning for Trading
    提到除以63日价格波动标准差作为 MACD 标准化,可以在波动性较高时平滑MACD
    这里在之后再做一个sigmoid变换将其归一化到 0~1
    :param X:
    :param k:
    :return:
    """
    assert 'diff' in X.columns.values
    assert 'dea' in X.columns.values
    assert 'close' in X.columns.values

    diff = X['diff'].values
    dea = X['dea'].values
    cdma = diff - dea
    close = X['close']
    close_yesterday = close.shift(1, axis=0)
    close_diff = close - close_yesterday
    for i in range(63):
        cdma[i] = np.NAN
    for i in range(63, len(close_diff) - 1):
        std = np.std(close_diff[i - 62:i + 1])
        cdma[i] = cdma[i] / std
    X['s_macd'] = cdma  # 1 / (1 + np.exp(-k * cdma))
    return X


def standardization_dma(X: pd.DataFrame, k=3) -> pd.DataFrame:
    """
    DMA 分为ddd线和ama线,由于ama线为ddd线的移动平均线，
     由于ddd线与macd计算方法基本一致，采取同样的归一化方法，
     由于ddd线的长线参数设置为50，略长于cdma所以增大了标准差的值
    :param X:
    :param k:
    :return:
    """
    assert 'ddd' in X.columns.values
    assert 'ama' in X.columns.values
    assert 'close' in X.columns.values

    ddd = X['ddd'].values
    ama = X['ama'].values
    dma = ddd - ama
    close = X['close']
    close_yesterday = close.shift(1, axis=0)
    close_diff = close - close_yesterday
    for i in range(89):
        dma[i] = np.NAN
    for i in range(89, len(close_diff) - 1):
        std = np.std(close[i - 88:i + 1])
        dma[i] = dma[i] / std
    X['s_dma'] = dma  # 1 / (1 + np.exp(-k * dma))
    return X


def standardization_mtm(X: pd.DataFrame, k=21.97) -> pd.DataFrame:
    """
    mtm 分为'mtm'与'mtmma'线，考虑到一般以两线交叉作为交易信号，
        所以对差值除以12日的移动平均线然后做sigmoid变换
        考虑到一般差值不超过0.1所以 k=14 可以保证-0.1~0.1的值被映射到80%的区间
    :param X:
    :param k:
    :return:
    """
    assert 'mtm' in X.columns.values
    assert 'mtmma' in X.columns.values
    assert 'close' in X.columns.values

    ma12 = np.squeeze(ta.MA(X, 12).values)
    mtm = X['mtm'].values
    mtmma = X['mtmma'].values
    ans = (mtm - mtmma) / ma12
    ans = 1 / (1 + np.exp(-k * ans))
    X['s_mtm'] = ans
    return X


def count_ma(origin_X: pd.DataFrame, *args):
    for tar in args:
        res = ta.MA(origin_X, tar)
        origin_X[f'ma_{tar}'] = res.values
    return origin_X


def set_standardization_feature(origin_X: pd.DataFrame) -> pd.DataFrame:
    """
        综合考虑，在数据处理时，将BBI、DMA、MACD、MTM除以当日的收盘价然后归一化，
        在 https://arxiv.org/abs/1911.10107
        Deep Reinforcement Learning for Trading
        也提到除以63日价格波动标准差作为 MACD 标准化
        则此时所有的指标范围均在-100~100左右， 除了CCI指标，KDJ的J线指标外
        @:param origin_X:
        @:return
    """
    origin_X = count_feature(origin_X)
    origin_X = count_ma(origin_X, *[5, 10, 20, 30])
    origin_X = open0(origin_X)
    origin_X = high0(origin_X)
    origin_X = low0(origin_X)

    origin_X = standardization_cci(origin_X)
    origin_X = standardization_mtm(origin_X)
    origin_X = standardization_macd(origin_X)
    origin_X = standardization_dma(origin_X)

    origin_X['s_bias'] = origin_X['bias'].apply(lambda x: 1 / (1 + np.exp(-0.1 * x)))

    col_need_div100 = ['k', 'd', 'j', 'mfi', 'priceosc', 'psy', 'roc', 'rocma', 'wr']
    for col in col_need_div100:
        new_col_name = f"s_{col}"
        origin_X[new_col_name] = origin_X[col].apply(lambda x: x / 100)

    origin_X['s_wr'] = origin_X['s_wr'].apply(lambda x: -x)
    origin_X['s_roc'] = origin_X['s_roc'].apply(lambda x: (x + 1) / 2)

    origin_X['s_rocma'] = origin_X['s_rocma'].apply(lambda x: (x + 1) / 2)
    # origin_X[['s_roc', 's_rocma']] = origin_X[['s_roc', 's_rocma']].apply(lambda x: 1 / (1 + np.exp(-20 * (x - 0.5))))

    origin_X['s_priceosc'] = origin_X['s_priceosc'].apply(lambda x: (x + 1) / 2)
    # origin_X['s_j'] = origin_X['s_j'].apply(lambda x: 1 / (1 + np.exp(-3 * (x - 0.5))))
    origin_X['s_change'] = origin_X['change'].apply(lambda x: 1 / (1 + np.exp(-45 * x)))
    return origin_X


def preprocess_sz50():
    import os
    from src.EnvironmentVariables import DATA_PATH, BASE_PATH
    from tqdm import tqdm
    from src.Utils.MyUtil import stock_adjust
    target_stocks = pd.read_csv(os.path.join(DATA_PATH, 'targetStocks.csv'))['id']
    dir_path = "E:/Codes/UndergraduateFinalDesign/all_stock_data/stock data"
    usecols = ["date",
               "open",
               "high",
               "low",
               "close",
               "volume",
               "turnover",
               "change",
               "adjust_price"]
    files_path = [os.path.join(dir_path, f"sh{f}.csv") for f in target_stocks]
    dataset = {}
    for i, stock_id in enumerate(tqdm(target_stocks, "Count:")):
        data = set_standardization_feature(
            stock_adjust(pd.read_csv(files_path[i],
                                     date_parser='date',
                                     usecols=usecols).sort_values('date')))
        dataset[stock_id] = data
    for key, val in tqdm(dataset.items(), "save:"):
        val.to_csv(os.path.join(BASE_PATH, 'data/preProcessedIndexData/sz50stock', f'sh{key}.csv'), index=False)


if __name__ == '__main__':
    preprocess_sz50()
